import numpy as np
from scipy.sparse import linalg, csr_matrix
import config
import math
import CoolProp as CP

class TwoFluidSolver:
    def __init__(self, subchannel):
        """初始化两流体模型求解器
        
        Args:
            subchannel: SubChannel对象，包含子通道的几何和物性信息
        """
        self.subchannel = subchannel
        
        # 网格参数
        self.n_axial = subchannel.n_axial
        self.n_channels = subchannel.n_channels
        self.dx = config.GEOMETRY_PARAMS['channel_length'] / (subchannel.n_axial - 1)
        self.dy = config.GEOMETRY_PARAMS['channel_width']
        
        # 时间步长控制
        self.dt = config.NUMERICAL_PARAMS['time_step']
        self.dt_min = config.NUMERICAL_PARAMS['min_time_step']
        self.dt_max = config.NUMERICAL_PARAMS['max_time_step']
        self.dt_increase = config.NUMERICAL_PARAMS['time_step_increase_factor']
        self.dt_decrease = config.NUMERICAL_PARAMS['time_step_decrease_factor']
        self.total_time = config.NUMERICAL_PARAMS['total_time']
        self.current_time = 0.0
        
        # 收敛控制
        self.max_iterations = config.NUMERICAL_PARAMS['max_iterations']
        self.convergence_tolerance = config.NUMERICAL_PARAMS['convergence_tolerance']
        
        # 松弛因子
        self.relaxation_factor = config.NUMERICAL_PARAMS['relaxation_factor']
        self.lateral_flow_under_relaxation = config.NUMERICAL_PARAMS['lateral_flow_under_relaxation']
        self.pressure_under_relaxation = config.NUMERICAL_PARAMS['pressure_under_relaxation']
        self.void_fraction_under_relaxation = config.NUMERICAL_PARAMS['void_fraction_under_relaxation']
        
        # 初始化存储上一步解的数组
        self._initialize_previous_solution()
        
        # 存储上一时间步的值用于收敛判断
        self.previous_values = {
            'void_fraction': np.zeros((self.n_axial, self.n_channels)),
            'pressure': np.zeros((self.n_axial, self.n_channels)),
            'vel_liquid': np.zeros((self.n_axial, self.n_channels)),
            'vel_vapor': np.zeros((self.n_axial, self.n_channels)),
            'temp_liquid': np.zeros((self.n_axial, self.n_channels)),
            'temp_vapor': np.zeros((self.n_axial, self.n_channels))
        }
        
        # 设置收敛判据
        self.convergence_criteria = {
            'void_fraction': 1e-4,
            'pressure': 1.0,  # Pa
            'velocity': 1e-3,  # m/s
            'temperature': 0.1  # K
        }
        
        # 设置松弛因子
        self.relaxation_factors = {
            'void_fraction': 0.3,
            'pressure': 0.3,
            'velocity': 0.3,
            'temperature': 0.3
        }
        
        # 初始化交错网格
        self._initialize_staggered_grid()
        
        # PKM求解器参数
        self.picard_tolerance = 0.01  # 1%的相对误差
        self.max_picard_iterations = 50
        self.min_dt = 1e-6  # 最小时间步长
        
        # PGMRES求解器参数
        self.gmres_tolerance = 1e-6
        self.max_gmres_iterations = 100
        self.restart = 30  # GMRES重启参数
        
    def _initialize_staggered_grid(self):
        """初始化交错网格系统"""
        # 标量网格（质量和能量）
        self.scalar_grid = {
            'pressure': np.zeros((self.n_axial, self.n_channels)),
            'void_fraction': np.zeros((self.n_axial, self.n_channels)),
            'enthalpy_liquid': np.zeros((self.n_axial, self.n_channels)),
            'enthalpy_vapor': np.zeros((self.n_axial, self.n_channels))
        }
        
        # 轴向向量网格（轴向动量）
        self.axial_grid = {
            'vel_liquid': np.zeros((self.n_axial + 1, self.n_channels)),
            'vel_vapor': np.zeros((self.n_axial + 1, self.n_channels))
        }
        
        # 横向向量网格（横向动量）
        self.lateral_grid = {
            'lateral_velocity': np.zeros((self.n_axial, self.n_channels + 1))
        }
        
    def _initialize_previous_solution(self):
        """初始化存储上一时间步解的数组"""
        self.previous_values = {
            'pressure': np.zeros((self.n_axial, self.n_channels)),
            'void_fraction': np.zeros((self.n_axial, self.n_channels)),
            'vel_liquid': np.zeros((self.n_axial, self.n_channels)),
            'vel_vapor': np.zeros((self.n_axial, self.n_channels)),
            'lateral_velocity': np.zeros((self.n_axial, self.n_channels)),
            'temp_liquid': np.zeros((self.n_axial, self.n_channels)),
            'temp_vapor': np.zeros((self.n_axial, self.n_channels)),
            'temp_wall': np.zeros((self.n_axial, self.n_channels))
        }
        
    def _store_previous_solution(self):
        """存储当前解为上一时间步的解"""
        self.previous_values['pressure'] = np.copy(self.subchannel.pressure)
        self.previous_values['void_fraction'] = np.copy(self.subchannel.void_fraction)
        self.previous_values['vel_liquid'] = np.copy(self.subchannel.vel_liquid)
        self.previous_values['vel_vapor'] = np.copy(self.subchannel.vel_vapor)
        self.previous_values['lateral_velocity'] = np.copy(self.subchannel.lateral_velocity)
        self.previous_values['temp_liquid'] = np.copy(self.subchannel.temp_liquid)
        self.previous_values['temp_vapor'] = np.copy(self.subchannel.temp_vapor)
        self.previous_values['temp_wall'] = np.copy(self.subchannel.temp_wall)
        
    def _get_previous_value(self, field_name, i, j):
        """获取上一时间步的值
        
        Args:
            field_name: 字段名称
            i, j: 网格位置
            
        Returns:
            上一时间步的值
        """
        if field_name not in self.previous_values:
            raise ValueError(f"Unknown field name: {field_name}")
        return self.previous_values[field_name][i,j]
        
    def solve_transient(self):
        """求解瞬态问题"""
        n_timesteps = int(self.total_time / self.dt)
        
        try:
            for step in range(n_timesteps):
                self.current_time = (step + 1) * self.dt
                print(f"\n开始计算时间步 {step+1}/{n_timesteps}, 当前时间: {self.current_time:.3f}s")
                print(f"当前时间步长: {self.dt:.2e}s")
                
                # 保存上一时间步的解
                self._store_previous_solution()
                
                # 求解当前时间步
                converged = self._solve_current_timestep()
                
                if not converged:
                    print("警告：当前时间步未收敛")
                    # 可以在这里添加时间步长调整的逻辑
                    
        except Exception as e:
            print(f"计算过程中出错: {str(e)}")
            raise
        
    def _solve_current_timestep(self):
        """求解当前时间步
        
        Returns:
            bool: 是否收敛
        """
        for iter in range(self.max_iterations):
            # 保存上一次迭代的解
            self._store_previous_solution()
            
            # 求解动量方程
            self._solve_momentum_equations()
            
            # 求解质量方程
            self._solve_mass_equations()
            
            # 求解能量方程
            self._solve_energy_equations()
            
            # 更新本构关系
            self._update_constitutive_relations()
            
            # 检查收敛性
            if self._check_convergence():
                print(f"迭代收敛，迭代次数: {iter+1}")
                return True
                
        print(f"警告：达到最大迭代次数 {self.max_iterations} 仍未收敛")
        return False
        
    def _check_convergence(self):
        """检查收敛性
        
        Returns:
            bool: 是否收敛
        """
        # 检查压力场的收敛性
        dp_max = np.max(np.abs(self.subchannel.pressure - self.previous_values['pressure']))
        if dp_max > self.convergence_tolerance:
            return False
            
        # 检查速度场的收敛性
        du_l_max = np.max(np.abs(self.subchannel.vel_liquid - self.previous_values['vel_liquid']))
        du_v_max = np.max(np.abs(self.subchannel.vel_vapor - self.previous_values['vel_vapor']))
        if max(du_l_max, du_v_max) > self.convergence_tolerance:
            return False
            
        # 检查空泡分数的收敛性
        da_max = np.max(np.abs(self.subchannel.void_fraction - self.previous_values['void_fraction']))
        if da_max > self.convergence_tolerance:
            return False
            
        # 检查温度场的收敛性
        dt_l_max = np.max(np.abs(self.subchannel.temp_liquid - self.previous_values['temp_liquid']))
        dt_v_max = np.max(np.abs(self.subchannel.temp_vapor - self.previous_values['temp_vapor']))
        if max(dt_l_max, dt_v_max) > self.convergence_tolerance:
            return False
            
        return True
        
    def _update_constitutive_relations(self):
        """更新本构关系"""
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                # 更新流体物性
                P = self.subchannel.pressure[i,j]
                T_l = self.subchannel.temp_liquid[i,j]
                T_v = self.subchannel.temp_vapor[i,j]
                
                # 更新基本物性
                self.subchannel.rho_liquid[i,j] = self._calculate_liquid_density(P, T_l)
                self.subchannel.rho_vapor[i,j] = self._calculate_vapor_density(P, T_v)
                self.subchannel.mu_liquid[i,j] = self._calculate_liquid_viscosity(T_l)
                self.subchannel.mu_vapor[i,j] = self._calculate_vapor_viscosity(T_v)
                self.subchannel.k_liquid[i,j] = self._calculate_liquid_conductivity(T_l)
                self.subchannel.k_vapor[i,j] = self._calculate_vapor_conductivity(T_v)
                self.subchannel.cp_liquid[i,j] = self._calculate_specific_heat(P, T_l, 'liquid')
                self.subchannel.cp_vapor[i,j] = self._calculate_specific_heat(P, T_v, 'vapor')
                
    def _calculate_liquid_density(self, P, T):
        """计算液相密度 (kg/m³)
        
        Args:
            P: 压力 (Pa)
            T: 温度 (K)
            
        Returns:
            液相密度 (kg/m³)
        """
        try:
            return CP.PropsSI('D', 'P', P, 'T', T, 'Water')
        except:
            # 使用简化的水密度计算公式
            T_C = T - 273.15
            return 1000 * (1 - ((T_C - 3.98)**2) / 50000)
            
    def _calculate_vapor_density(self, P, T):
        """计算气相密度 (kg/m³)
        
        Args:
            P: 压力 (Pa)
            T: 温度 (K)
            
        Returns:
            气相密度 (kg/m³)
        """
        try:
            return CP.PropsSI('D', 'P', P, 'T', T, 'Water')
        except:
            # 使用理想气体状态方程
            R = 461.5  # 水蒸气气体常数 J/(kg·K)
            return P / (R * T)
            
    def _calculate_liquid_viscosity(self, T):
        """计算液相粘度 (Pa·s)
        
        Args:
            T: 温度 (K)
            
        Returns:
            液相粘度 (Pa·s)
        """
        try:
            P_sat = CP.PropsSI('P', 'T', T, 'Q', 0, 'Water')
            return CP.PropsSI('V', 'P', P_sat, 'T', T, 'Water')
        except:
            # 使用简化的水粘度计算公式
            T_C = T - 273.15
            return 0.001 * (1 + 0.0337 * T_C + 0.000221 * T_C**2) ** (-1)
            
    def _calculate_vapor_viscosity(self, T):
        """计算气相粘度 (Pa·s)
        
        Args:
            T: 温度 (K)
            
        Returns:
            气相粘度 (Pa·s)
        """
        try:
            P_sat = CP.PropsSI('P', 'T', T, 'Q', 1, 'Water')
            return CP.PropsSI('V', 'P', P_sat, 'T', T, 'Water')
        except:
            # 使用Sutherland公式
            mu_0 = 1.85e-5
            T_0 = 373.15
            S = 620
            return mu_0 * (T/T_0)**1.5 * (T_0 + S)/(T + S)
            
    def _calculate_liquid_conductivity(self, T):
        """计算液相导热系数 (W/(m·K))
        
        Args:
            T: 温度 (K)
            
        Returns:
            液相导热系数 (W/(m·K))
        """
        try:
            P_sat = CP.PropsSI('P', 'T', T, 'Q', 0, 'Water')
            return CP.PropsSI('L', 'P', P_sat, 'T', T, 'Water')
        except:
            # 使用简化的水导热系数计算公式
            T_C = T - 273.15
            return 0.6 - 0.0005 * T_C
            
    def _calculate_vapor_conductivity(self, T):
        """计算气相导热系数 (W/(m·K))
        
        Args:
            T: 温度 (K)
            
        Returns:
            气相导热系数 (W/(m·K))
        """
        try:
            P_sat = CP.PropsSI('P', 'T', T, 'Q', 1, 'Water')
            return CP.PropsSI('L', 'P', P_sat, 'T', T, 'Water')
        except:
            # 使用简化的水蒸气导热系数计算公式
            return 0.016 + 5e-5 * (T - 273.15)
            
    def _calculate_specific_heat(self, P, T, phase):
        """计算比热容 (J/(kg·K))
        
        Args:
            P: 压力 (Pa)
            T: 温度 (K)
            phase: 相态 ('liquid' 或 'vapor')
            
        Returns:
            比热容 (J/(kg·K))
        """
        try:
            if phase == 'liquid':
                return CP.PropsSI('C', 'P', P, 'T', T, 'Water')
            else:  # vapor
                return CP.PropsSI('C', 'P', P, 'T', T, 'Water')
        except:
            # 如果CoolProp计算失败，使用近似值
            return 4200 if phase == 'liquid' else 2000
        
    def _solve_momentum_equations(self):
        """求解动量方程"""
        # 轴向动量方程
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                self._solve_axial_momentum('l', i, j)  # 液相
                self._solve_axial_momentum('v', i, j)  # 气相
                
        # 横向动量方程
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                self._solve_lateral_momentum('l', i, j)  # 液相
                self._solve_lateral_momentum('v', i, j)  # 气相
                
    def _solve_axial_momentum(self, k, i, j):
        """求解轴向动量方程
        eq.2: ∂(αₖρₖuₖ)/∂t + ∂(αₖρₖuₖuₖ)/∂x + ∂(αₖρₖuₖvₖ)/∂y = -αₖ∂P/∂x - αₖρₖg - Fxwₖ + Flv + uΓₖ
        """
        # 获取流体状态
        if k == 'l':
            alpha = 1 - self.subchannel.void_fraction[i,j]
            rho = self.subchannel.rho_liquid[i,j]
            u = self.subchannel.vel_liquid[i,j]
            v = self.subchannel.lateral_velocity[i,j]
        else:
            alpha = self.subchannel.void_fraction[i,j]
            rho = self.subchannel.rho_vapor[i,j]
            u = self.subchannel.vel_vapor[i,j]
            v = self.subchannel.lateral_velocity[i,j]
            
        # 计算压力梯度
        if i == 0:
            dP_dx = (self.subchannel.pressure[i+1,j] - self.subchannel.pressure[i,j]) / self.subchannel.dx
        elif i == self.n_axial - 1:
            dP_dx = (self.subchannel.pressure[i,j] - self.subchannel.pressure[i-1,j]) / self.subchannel.dx
        else:
            dP_dx = (self.subchannel.pressure[i+1,j] - self.subchannel.pressure[i-1,j]) / (2 * self.subchannel.dx)
            
        # 计算各种力
        f_wall = self._calculate_wall_friction(k, i, j)  # 壁面摩擦力
        f_int = self._calculate_interfacial_friction(k, i, j)  # 界面剪切力
        f_vm = self._calculate_virtual_mass_force(k, i, j)  # 虚拟质量力
        
        # 计算相变动量源项
        gamma = self._calculate_phase_change(k, i, j)
        mom_phase = u * gamma
        
        # 计算动量方程右端项
        R = -alpha * dP_dx - alpha * rho * config.PHYSICAL_PARAMS['gravity'] - f_wall + f_int + mom_phase + f_vm
        
        # 计算速度变化
        du_dt = R / (alpha * rho)
        
        # 更新速度
        if k == 'l':
            self.subchannel.vel_liquid[i,j] += du_dt * self.dt
        else:
            self.subchannel.vel_vapor[i,j] += du_dt * self.dt
            
    def _solve_lateral_momentum(self, k, i, j):
        """求解横向动量方程
        eq.3: ∂(αₖρₖvₖ)/∂t + ∂(αₖρₖvₖuₖ)/∂x + ∂(αₖρₖvₖvₖ)/∂y = -αₖ∂P/∂y - Fywₖ + Flv - vΓₖ + TyMₖ
        """
        # 获取相邻通道
        connected_channels = self.subchannel.get_connected_channels(j)
        
        for neighbor_j, conn_type, conn_width in connected_channels:
            if conn_type == 1:  # 只处理子通道间连接
                # 获取流体状态
                if k == 'l':
                    alpha = 1 - self.subchannel.void_fraction[i,j]
                    rho = self.subchannel.rho_liquid[i,j]
                    v = self.subchannel.lateral_velocity[i,j]
                else:
                    alpha = self.subchannel.void_fraction[i,j]
                    rho = self.subchannel.rho_vapor[i,j]
                    v = self.subchannel.lateral_velocity[i,j]
                    
                # 计算横向压力梯度
                dP_dy = (self.subchannel.pressure[i,neighbor_j] - self.subchannel.pressure[i,j]) / conn_width
                
                # 计算各种力
                f_wall_y = self._calculate_lateral_wall_friction(k, i, j, neighbor_j)
                f_int_y = self._calculate_lateral_interfacial_friction(k, i, j, neighbor_j)
                
                # 计算相变和湍流混合项
                gamma = self._calculate_phase_change(k, i, j)
                T_yM = self._calculate_turbulent_mixing(k, i, j, neighbor_j)
                
                # 计算横向动量方程右端项
                R = -alpha * dP_dy - f_wall_y + f_int_y - v * gamma + T_yM
                
                # 计算横向速度变化
                dv_dt = R / (alpha * rho)
                
                # 更新横向速度
                v_new = v + dv_dt * self.dt
                
                # 使用松弛因子更新
                relax = config.NUMERICAL_PARAMS['lateral_flow_under_relaxation']
                self.subchannel.lateral_velocity[i,j] = v_new * relax + v * (1 - relax)
                
    def _solve_energy_equations(self):
        """求解能量方程
        eq.4: ∂(αₖρₖhₖ)/∂t + ∂(αₖρₖuₖhₖ)/∂x + ∂(αₖρₖvₖhₖ)/∂y = αₖ∂P/∂t + qwₖ + Γₖh + Tke
        """
        for i in range(self.n_axial):
            for j in range(self.n_channels):
                # 计算液相和气相的能量方程
                for k in ['l', 'v']:
                    # 获取流体状态
                    if k == 'l':
                        alpha = 1 - self.subchannel.void_fraction[i,j]
                        rho = self.subchannel.rho_liquid[i,j]
                        h = self.subchannel.enthalpy_liquid[i,j]
                        u = self.subchannel.vel_liquid[i,j]
                        v = self.subchannel.lateral_velocity[i,j]
                    else:
                        alpha = self.subchannel.void_fraction[i,j]
                        rho = self.subchannel.rho_vapor[i,j]
                        h = self.subchannel.enthalpy_vapor[i,j]
                        u = self.subchannel.vel_vapor[i,j]
                        v = self.subchannel.lateral_velocity[i,j]
                        
                    # 计算轴向能量通量导数
                    if i == 0:
                        dE_dx = (alpha * rho * u * h - self._get_inlet_energy_flux(k, j)) / self.subchannel.dx
                    else:
                        dE_dx = (alpha * rho * u * h - 
                                self._get_energy_flux(k, i-1, j)) / self.subchannel.dx
                    
                    # 计算横向能量通量
                    Q = self._calculate_lateral_energy_flux(k, i, j)
                    
                    # 计算压力功
                    dP_dt = (self.subchannel.pressure[i,j] - self.pressure_old[i,j]) / self.dt
                    W_p = alpha * dP_dt
                    
                    # 计算壁面传热
                    q_wall = self._calculate_wall_heat_transfer(k, i, j)
                    
                    # 计算相变热
                    gamma = self._calculate_phase_change(k, i, j)
                    h_sat = self._calculate_saturation_enthalpy(self.subchannel.pressure[i,j], k)
                    Q_phase = gamma * h_sat
                    
                    # 计算湍流混合能量
                    T_ke = self._calculate_turbulent_energy_mixing(k, i, j)
                    
                    # 计算能量方程右端项
                    R = W_p + q_wall + Q_phase + T_ke - dE_dx - Q
                    
                    # 计算温度变化
                    if k == 'l':
                        dT_dt = R / (alpha * rho * self.subchannel.cp_liquid[i,j])
                        self.subchannel.temp_liquid[i,j] += dT_dt * self.dt
                    else:
                        dT_dt = R / (alpha * rho * self.subchannel.cp_vapor[i,j])
                        self.subchannel.temp_vapor[i,j] += dT_dt * self.dt
        
    def _calculate_axial_energy_flux(self, k, i, j):
        """计算轴向能量通量
        
        Args:
            k: 相态 ('l' 或 'v')
            i: 轴向索引
            j: 子通道索引
            
        Returns:
            轴向能量通量 (W/m²)
        """
        if k == 'l':
            rho = self.subchannel.rho_liquid[i,j]
            v = self.subchannel.vel_liquid[i,j]
            h = self.subchannel.enthalpy_liquid[i,j]
            alpha = max(1e-6, 1 - self.subchannel.void_fraction[i,j])
        else:
            rho = self.subchannel.rho_vapor[i,j]
            v = self.subchannel.vel_vapor[i,j]
            h = self.subchannel.enthalpy_vapor[i,j]
            alpha = max(1e-6, self.subchannel.void_fraction[i,j])
            
        return rho * alpha * v * h
        
    def _calculate_lateral_energy_flux(self, k, i, j):
        """计算横向能量通量
        
        Args:
            k: 相态 ('l' 或 'v')
            i: 轴向索引
            j: 子通道索引
            
        Returns:
            横向能量通量 (W/m²)
        """
        # 获取相邻通道
        connected_channels = self.subchannel.get_connected_channels(j)
        Q_net = 0.0
        
        for neighbor_j, conn_type, conn_width in connected_channels:
            if conn_type == 1:  # 只处理子通道间连接
                # 获取流体状态
                if k == 'l':
                    rho = self.subchannel.rho_liquid[i,j]
                    h = self.subchannel.enthalpy_liquid[i,j]
                    h_n = self.subchannel.enthalpy_liquid[i,neighbor_j]
                    alpha = max(1e-6, 1 - self.subchannel.void_fraction[i,j])
                else:
                    rho = self.subchannel.rho_vapor[i,j]
                    h = self.subchannel.enthalpy_vapor[i,j]
                    h_n = self.subchannel.enthalpy_vapor[i,neighbor_j]
                    alpha = max(1e-6, self.subchannel.void_fraction[i,j])
                    
                # 计算横向度
                w = self.subchannel.lateral_velocity[i,j]
                
                # 计算对流能量通量
                Q_conv = rho * alpha * w * h * conn_width
                
                # 计算湍流混合能量通量
                beta = config.LATERAL_FLOW_PARAMS['mixing_coefficient']
                Q_mix = beta * rho * alpha * abs(w) * (h_n - h) * conn_width
                
                # 累加净能量通量
                Q_net += Q_conv + Q_mix
                
        return Q_net
        
    def _calculate_energy_residual(self, k, i, j, E, Q, q_wall, q_i, gamma, h_fg):
        """计算能量方程残差
        
        Args:
            k: 相态 ('l' 或 'v')
            i: 轴向索引
            j: 子通道索引
            E: 轴向能量通量 (W/m²)
            Q: 横向能量通量 (W/m²)
            q_wall: 壁面传热热流密度 (W/m²)
            q_i: 界面传热热流密度 (W/m²)
            gamma: 相变率 (kg/m³s)
            h_fg: 汽化潜热 (J/kg)
            
        Returns:
            能量方程残差 (W/m³)
        """
        # 计算轴向导数
        if i == 0:  # 入口边界
            dE_dz = (E - self._get_inlet_energy_flux(k, j)) / self.subchannel.dx
        elif i == self.n_axial - 1:  # 出口边界
            dE_dz = (self._get_outlet_energy_flux(k, j) - E) / self.subchannel.dx
        else:  # 内部节点
            dE_dz = (E - self._get_energy_flux(k, i-1, j)) / self.subchannel.dx
            
        # 计算能量方程残差
        if k == 'l':
            R_e = dE_dz + Q - q_wall - q_i - gamma * h_fg
        else:
            R_e = dE_dz + Q - q_wall - q_i + gamma * h_fg
            
        return R_e
        
    def _update_temperatures(self, i, j, R_e_l, R_e_v):
        """更新温度场
        
        Args:
            i: 轴向索引
            j: 子通道索引
            R_e_l: 液相能量方程残差 (W/m³)
            R_e_v: 气相能量方程残差 (W/m³)
        """
        # 获取流体状态
        alpha = self.subchannel.void_fraction[i,j]
        rho_l = self.subchannel.rho_liquid[i,j]
        rho_v = self.subchannel.rho_vapor[i,j]
        cp_l = self.subchannel.cp_liquid[i,j]
        cp_v = self.subchannel.cp_vapor[i,j]
        
        # 计算温度变化率
        dT_l = -self.dt * R_e_l / (rho_l * (1 - alpha) * cp_l)
        dT_v = -self.dt * R_e_v / (rho_v * alpha * cp_v)
        
        # 更新温度
        T_l_old = self.subchannel.temp_liquid[i,j]
        T_v_old = self.subchannel.temp_vapor[i,j]
        
        T_l_new = T_l_old + dT_l
        T_v_new = T_v_old + dT_v
        
        # 限制温度变化幅度
        max_dT = 10.0  # 最大允许温度变化 (K)
        T_l_new = T_l_old + max(-max_dT, min(max_dT, dT_l))
        T_v_new = T_v_old + max(-max_dT, min(max_dT, dT_v))
        
        # 使用松弛因子更新温度
        relax = config.NUMERICAL_PARAMS.get('temperature_under_relaxation', 0.8)
        self.subchannel.temp_liquid[i,j] = T_l_new * relax + T_l_old * (1 - relax)
        self.subchannel.temp_vapor[i,j] = T_v_new * relax + T_v_old * (1 - relax)
        
    def _get_inlet_energy_flux(self, k, j):
        """获取入口能量通量
        
        Args:
            k: 相态 ('l' 或 'v')
            j: 子通道索引
            
        Returns:
            入口能量通量 (W/m²)
        """
        # 从初始条件获取
        T_inlet = config.INITIAL_CONDITIONS['inlet_temperature']
        G_inlet = config.INITIAL_CONDITIONS['inlet_mass_flux']
        alpha_inlet = config.INITIAL_CONDITIONS['inlet_void_fraction']
        
        if k == 'l':
            h = self._calculate_liquid_enthalpy(T_inlet)
            return G_inlet * (1 - alpha_inlet) * h
        else:
            h = self._calculate_vapor_enthalpy(config.INITIAL_CONDITIONS['inlet_pressure'], T_inlet)
            return G_inlet * alpha_inlet * h
            
    def _get_outlet_energy_flux(self, k, i, j):
        """获取出口能量通量（简单外推）
        
        Args:
            k: 相态 ('l' 或 'v')
            i: 轴向索引
            j: 子通道索引
            
        Returns:
            出口能量通量 (W/m²)
        """
        return self._get_energy_flux(k, i, j)
        
    def _get_energy_flux(self, k, i, j):
        """获取指定位置的能量通量
        
        Args:
            k: 相态 ('l' 或 'v')
            i: 轴向索引
            j: 子通道索引
            
        Returns:
            能量通量 (W/m²)
        """
        if k == 'l':
            return self.subchannel.rho_liquid[i,j] * (1 - self.subchannel.void_fraction[i,j]) * \
                   self.subchannel.vel_liquid[i,j] * self.subchannel.enthalpy_liquid[i,j]
        else:
            return self.subchannel.rho_vapor[i,j] * self.subchannel.void_fraction[i,j] * \
                   self.subchannel.vel_vapor[i,j] * self.subchannel.enthalpy_vapor[i,j]
        
    def _calculate_wall_friction(self, k, i, j):
        """计算轴向壁面摩擦力
        
        Args:
            k: 相态 ('l' 或 'v')
            i: 轴向索引
            j: 子通道索引
            
        Returns:
            轴向面摩擦力 (N/m³)
        """
        # 获取流体参数
        if k == 'l':
            alpha = max(1e-6, 1 - self.subchannel.void_fraction[i,j])
            rho = self.subchannel.rho_liquid[i,j]
            mu = self._calculate_liquid_viscosity(self.subchannel.temp_liquid[i,j])
            v = self.subchannel.vel_liquid[i,j]
        else:
            alpha = max(1e-6, self.subchannel.void_fraction[i,j])
            rho = self.subchannel.rho_vapor[i,j]
            mu = self._calculate_vapor_viscosity(self.subchannel.temp_vapor[i,j])
            v = self.subchannel.vel_vapor[i,j]
            
        # 计算雷诺数
        D_h = self.subchannel.hydraulic_diameter[j]
        Re = abs(rho * v * D_h / mu) if mu > 0 else 1e6
        
        # 计算摩擦系数
        f = self._calculate_friction_factor(Re)
        
        # 计算壁面剪切应力
        tau_w = 0.5 * f * rho * v * abs(v)
        
        # 考虑两相流动的影响
        phi2 = self._calculate_two_phase_friction_multiplier(k, i, j)
        tau_w *= phi2
        
        # 计算壁面摩擦力
        P_w = self.subchannel.wetted_perimeter[j]  # 湿周
        A_c = self.subchannel.flow_area[j]  # 流通面积
        
        return tau_w * P_w / (A_c * alpha)
        
    def _calculate_friction_factor(self, Re):
        """计算摩擦系数
        
        Args:
            Re: 雷诺数
            
        Returns:
            摩擦系数
        """
        if Re < 2300:  # 层流
            return 64 / max(1.0, Re)
        else:  # 湍流，使用Blasius公式
            return 0.316 / max(1.0, Re)**0.25
            
    def _calculate_two_phase_friction_multiplier(self, k, i, j):
        """计算两相流动摩擦系数修正因子
        
        Args:
            k: 相态 ('l' 或 'v')
            i: 轴向索引
            j: 子通道索引
            
        Returns:
            两相流动摩擦系数修正因子
        """
        alpha = self.subchannel.void_fraction[i,j]
        
        if alpha < 1e-6:  # 纯液相
            return 1.0 if k == 'l' else 0.0
        elif alpha > 0.9999:  # 纯气相
            return 0.0 if k == 'l' else 1.0
            
        # 计算Martinelli参数
        rho_l = self.subchannel.rho_liquid[i,j]
        rho_v = self.subchannel.rho_vapor[i,j]
        mu_l = self._calculate_liquid_viscosity(self.subchannel.temp_liquid[i,j])
        mu_v = self._calculate_vapor_viscosity(self.subchannel.temp_vapor[i,j])
        
        # 防止除零
        mu_l = max(1e-6, mu_l)
        mu_v = max(1e-6, mu_v)
        
        X = ((1 - alpha) / alpha)**0.9 * \
            (rho_v / rho_l)**0.5 * \
            (mu_l / mu_v)**0.1
            
        # 使用Chisholm相关式
        if k == 'l':
            return 1 + 20/X + 1/X**2
        else:
            return X**2 * (1 + 20/X + 1/X**2)
            
    def _calculate_interfacial_friction(self, k, i, j):
        """计算界面摩擦力
        
        Args:
            k: 相态 ('l' 或 'v')
            i: 轴向索引
            j: 子通道索引
            
        Returns:
            界面摩擦力 (N/m³)
        """
        # 获取流体状态
        alpha = self.subchannel.void_fraction[i,j]
        rho_l = self.subchannel.rho_liquid[i,j]
        v_l = self.subchannel.vel_liquid[i,j]
        v_v = self.subchannel.vel_vapor[i,j]
        
        # 计算相对速度
        v_r = v_v - v_l
        
        # 计算界面面积密度
        a_i = self._calculate_interfacial_area_density(i, j)
        
        # 计算界面阻力系数
        C_i = self._calculate_interfacial_friction_coefficient(i, j)
        
        # 计算界面摩擦力
        f_i = 0.5 * C_i * rho_l * a_i * v_r * abs(v_r)
        
        # 根据相态返回力的方向
        if k == 'l':
            return f_i
        else:
            return -f_i
            
    def _calculate_interfacial_friction_coefficient(self, i, j):
        """计算界面阻力系数
        
        Args:
            i: 轴向索引
            j: 子通道索引
            
        Returns:
            界面阻力系数
        """
        # 获取流体状态
        alpha = self.subchannel.void_fraction[i,j]
        rho_l = self.subchannel.rho_liquid[i,j]
        mu_l = self._calculate_liquid_viscosity(self.subchannel.temp_liquid[i,j])
        v_r = abs(self.subchannel.vel_vapor[i,j] - self.subchannel.vel_liquid[i,j])
        
        # 计算气泡雷诺数
        D_b = max(0.001, self.subchannel.hydraulic_diameter[j] * 0.1)  # 气泡直径
        Re_b = rho_l * v_r * D_b / max(1e-6, mu_l)
        
        # 使用标准阻力系数关联式
        if Re_b < 1:
            C_i = 24 / max(1.0, Re_b)
        elif Re_b < 1000:
            C_i = 24 / max(1.0, Re_b) * (1 + 0.15 * Re_b**0.687)
        else:
            C_i = 0.44
            
        return C_i
        
    def _calculate_heat_transfer_coefficient(self, k, i, j):
        """计算对流换热系数
        
        Args:
            k: 相态 ('l' 或 'v')
            i: 轴向索引
            j: 子通道索引
            
        Returns:
            传热系数 (W/m²K)
        """
        # 获取温度
        T_w = self.subchannel.temp_wall[i,j]
        T_sat = self._calculate_saturation_temperature(self.subchannel.pressure[i,j])
        T_l = self.subchannel.temp_liquid[i,j]
        P = self.subchannel.pressure[i,j]
        alpha = self.subchannel.void_fraction[i,j]
        
        if k == 'l':  # 液相
            # 判断沸腾状态
            if T_w < T_sat:  # 单相液态
                # 使用Dittus-Boelter相关式
                Re = self.subchannel.rho_liquid[i,j] * abs(self.subchannel.vel_liquid[i,j]) * \
                     self.subchannel.hydraulic_diameter[j] / self._calculate_liquid_viscosity(T_l)
                Pr = self.subchannel.cp_liquid[i,j] * self._calculate_liquid_viscosity(T_l) / \
                     self.subchannel.k_liquid[i,j]
                Nu = 0.023 * Re**0.8 * Pr**0.4
                h_l = Nu * self.subchannel.k_liquid[i,j] / self.subchannel.hydraulic_diameter[j]
                
            elif T_l < T_sat:  # 过冷核态沸腾
                # 使用Thom相关式结合冷凝修正
                F_gam = self._calculate_condensation_factor(i, j)
                h_l = (1 - F_gam) * self._calculate_thom_correlation(P, T_w, T_sat, T_l)
                
            else:  # 饱和核态沸腾
                # 使用Thom相关式
                h_l = self._calculate_thom_correlation(P, T_w, T_sat, T_l)
                
            return h_l
            
        else:  # 气相
            if alpha < 0.999:  # 非单相气态
                return 0  # 忽略壁面与气泡之间的传热
            else:  # 单相气态
                # 使用Dittus-Boelter相关式
                Re = self.subchannel.rho_vapor[i,j] * abs(self.subchannel.vel_vapor[i,j]) * \
                     self.subchannel.hydraulic_diameter[j] / self._calculate_vapor_viscosity(T_l)
                Pr = self.subchannel.cp_vapor[i,j] * self._calculate_vapor_viscosity(T_l) / \
                     self.subchannel.k_vapor[i,j]
                Nu = 0.023 * Re**0.8 * Pr**0.4
                h_v = Nu * self.subchannel.k_vapor[i,j] / self.subchannel.hydraulic_diameter[j]
                return h_v
                
    def _calculate_thom_correlation(self, P, T_w, T_sat, T_l):
        """计算Thom相关式的传热系数
        
        Args:
            P: 压力 (Pa)
            T_w: 壁面温度 (K)
            T_sat: 饱和温度 (K)
            T_l: 液相温度 (K)
            
        Returns:
            传热系数 (W/m²K)
        """
        # 将压力转换为bar
        P_bar = P / 1e5
        
        # 计算Thom相关式
        h = np.exp(2 * P_bar / 1260 * (T_w - T_sat) / (2722 * (T_w - T_l)))
        
        return h
        
    def _calculate_condensation_factor(self, i, j):
        """计算气泡冷凝系数
        
        Args:
            i: 轴向索引
            j: 子通道索引
            
        Returns:
            冷凝系数
        """
        # 获取流体状态
        T_l = self.subchannel.temp_liquid[i,j]
        T_sat = self._calculate_saturation_temperature(self.subchannel.pressure[i,j])
        
        # 计算过冷度
        dT_sub = T_sat - T_l
        
        # 简化的冷凝系数型
        F_gam = min(1.0, max(0.0, dT_sub / 10))  # 假设10K为完全冷凝的过冷度
        
        return F_gam
        
    def _calculate_phase_change(self, k, i, j):
        """计算相变质量流量
        
        Args:
            k: 相态 ('l' 或 'v')
            i: 轴向索引
            j: 子通道索引
            
        Returns:
            相变质量流量 (kg/m³s)
        """
        # 获取流体状态
        P = self.subchannel.pressure[i,j]
        T_l = self.subchannel.temp_liquid[i,j]
        T_v = self.subchannel.temp_vapor[i,j]
        T_sat = self._calculate_saturation_temperature(P)
        
        # 获取焓值
        h_l = self.subchannel.enthalpy_liquid[i,j]
        h_v = self.subchannel.enthalpy_vapor[i,j]
        h_l_sat = self._calculate_saturation_enthalpy(P, 'liquid')
        h_v_sat = self._calculate_saturation_enthalpy(P, 'vapor')
        
        # 计算各种状态下的相变率
        Gamma_l_sh = 0  # 液相过热
        Gamma_v_sh = 0  # 气相过热
        Gamma_l_sc = 0  # 液相过冷
        Gamma_v_sc = 0  # 气相过冷
        
        # 液相过热/过冷判断
        if T_l > T_sat:  # 液相过热
            C_l_sh = 0.1  # 液相过热传热系数
            Gamma_l_sh = C_l_sh * self.subchannel.rho_liquid[i,j] * \
                        self.subchannel.cp_liquid[i,j] * (h_l - h_l_sat)
        elif T_l < T_sat:  # 液相过冷
            C_l_sc = 0.1  # 液相过冷传热系数
            Gamma_l_sc = C_l_sc * self.subchannel.rho_liquid[i,j] * \
                        self.subchannel.cp_liquid[i,j] * (h_l_sat - h_l)
                        
        # 气相过热/过冷判断
        if T_v > T_sat:  # 气相过热
            C_v_sh = 0.1  # 气相过热传热系数
            Gamma_v_sh = C_v_sh * self.subchannel.rho_vapor[i,j] * \
                        self.subchannel.cp_vapor[i,j] * (h_v - h_v_sat)
        elif T_v < T_sat:  # 气相过冷
            C_v_sc = 0.1  # 气相过冷传热系数
            Gamma_v_sc = C_v_sc * self.subchannel.rho_vapor[i,j] * \
                        self.subchannel.cp_vapor[i,j] * (h_v_sat - h_v)
                        
        # 计算净相变率
        Gamma = Gamma_l_sh + Gamma_v_sh - Gamma_l_sc - Gamma_v_sc
        
        # 根据相态返回相应的相变率
        if k == 'l':
            return -Gamma  # 液相减少
        else:
            return Gamma  # 气相增加
            
    def _calculate_turbulent_mixing(self, k, i, j, neighbor_j):
        """计算湍流混合和空泡漂移质量流量
        
        Args:
            k: 相态 ('l' 或 'v')
            i: 轴向索引
            j: 当前通道索引
            neighbor_j: 相邻通道索引
            
        Returns:
            湍流混合质量流量 (kg/m²s)
        """
        # 获取几何参数
        L_gap = config.GEOMETRY_PARAMS['gap_width']
        L_tur = config.GEOMETRY_PARAMS['mixing_length']
        
        # 获取流体状态
        rho_l = self.subchannel.rho_liquid[i,j]
        rho_v = self.subchannel.rho_vapor[i,j]
        alpha_i = self.subchannel.void_fraction[i,j]
        alpha_j = self.subchannel.void_fraction[i,neighbor_j]
        
        # 计算平衡空泡分布（简化模型）
        alpha_i_equil = (alpha_i + alpha_j) / 2
        
        # 计算涡流扩散率
        C_e = config.LATERAL_FLOW_PARAMS['mixing_coefficient']
        
        # 计算湍流混合质量流量
        T_m = C_e * L_gap * L_tur * (rho_l - rho_v) * \
              ((alpha_j - alpha_i) - (alpha_j - alpha_i_equil))
              
        return T_m
        
    def solve_step(self):
        """使用PKM方法求解一个时间步"""
        dt = self.dt
        converged = False
        
        while not converged and dt >= self.min_dt:
            try:
                converged = self._picard_iteration(dt)
                if not converged:
                    dt *= 0.5
                    print(f"Reducing time step to {dt}")
            except np.linalg.LinAlgError:
                dt *= 0.5
                print(f"Linear solver failed. Reducing time step to {dt}")
                if dt < self.min_dt:
                    raise RuntimeError("Failed to converge even with minimum time step")
                    
        return converged, dt
        
    def _picard_iteration(self, dt):
        """执行Picard迭代"""
        for iter in range(self.max_picard_iterations):
            # 1. 计算加热棒与流场之间的热通量
            self._calculate_heat_flux()
            
            # 2. 求解动量方程获得质量流量
            self._solve_momentum_equations()
            
            # 3. 构建并求解质量-能量方程组
            residual = self._solve_mass_energy_equations()
            
            # 4. 更新两相本构关系
            self._update_constitutive_relations()
            
            # 检查收敛性
            if residual < self.picard_tolerance:
                return True
                
        return False
        
    def _solve_momentum_equations(self):
        """求解动量方程"""
        # 轴向动量方程（交错网格）
        for i in range(self.n_axial + 1):
            for j in range(self.n_channels):
                for k in ['l', 'v']:
                    self._solve_axial_momentum(k, i, j)
                    
        # 横向动量方程（交错网格）
        for i in range(self.n_axial):
            for j in range(self.n_channels + 1):
                self._solve_lateral_momentum(i, j)
                
    def _solve_mass_energy_equations(self):
        """使用PGMRES求解质量-能量方程组"""
        # 构建系统矩阵和右端向量
        A = self._build_mass_energy_matrix()
        b = self._calculate_mass_energy_residual()
        
        # 构建预处理器（仅考虑轴向耦合）
        M = self._build_axial_preconditioner(A)
        
        # 使用PGMRES求解
        x, info = sparse.linalg.gmres(
            A, b,
            M=M,
            tol=self.gmres_tolerance,
            maxiter=self.max_gmres_iterations,
            restart=self.restart
        )
        
        if info != 0:
            raise np.linalg.LinAlgError("GMRES failed to converge")
            
        # 更新解
        self._update_solution(x)
        
        return np.linalg.norm(A @ x - b) / np.linalg.norm(b)
        
    def _build_axial_preconditioner(self, A):
        """构建基于物理过程的预处理器"""
        # 提取主对角线和轴向耦合项
        n = A.shape[0]
        M = sparse.lil_matrix((n, n))
        
        for i in range(n):
            # 主对角线
            M[i,i] = A[i,i]
            
            # 轴向上下相邻元素
            if i >= self.n_channels:  # 有上方元素
                M[i,i-self.n_channels] = A[i,i-self.n_channels]
            if i < n - self.n_channels:  # 有下方元素
                M[i,i+self.n_channels] = A[i,i+self.n_channels]
                
        # 转换为CSR格式以提高求解效率
        M = M.tocsr()
        
        # 使用ILU分解作为预处理器
        return sparse.linalg.spilu(M)
        
    def _solve_axial_momentum(self, k, i, j):
        """在交错网格上求解轴向动量方程"""
        if k == 'l':
            alpha = 1 - self._interpolate_void_fraction(i, j)
            rho = self.subchannel.rho_liquid[i,j]
            u = self.axial_grid['vel_liquid'][i,j]
        else:
            alpha = self._interpolate_void_fraction(i, j)
            rho = self.subchannel.rho_vapor[i,j]
            u = self.axial_grid['vel_vapor'][i,j]
            
        # 计算动量通量（使用迎风格式）
        flux = self._calculate_momentum_flux(k, i, j)
        
        # 计算压力梯度
        dp_dx = self._calculate_pressure_gradient(i, j)
        
        # 计算源项
        source = self._calculate_momentum_source(k, i, j)
        
        # 更新速度
        u_new = u + (source - flux - alpha * dp_dx) * self.dt / (alpha * rho)
        
        if k == 'l':
            self.axial_grid['vel_liquid'][i,j] = u_new
        else:
            self.axial_grid['vel_vapor'][i,j] = u_new
            
    def _solve_lateral_momentum(self, i, j):
        """在交错网格上求解横向动量方程"""
        # 获取相邻通道的压力差
        dp_dy = self._calculate_lateral_pressure_gradient(i, j)
        
        # 计算横向动量通量
        flux = self._calculate_lateral_momentum_flux(i, j)
        
        # 计算源项（包括湍流混合和空泡漂移）
        source = self._calculate_lateral_momentum_source(i, j)
        
        # 更新横向速度
        v_new = (self.lateral_grid['lateral_velocity'][i,j] + 
                (source - flux - dp_dy) * self.dt)
                
        self.lateral_grid['lateral_velocity'][i,j] = v_new
        
    def _interpolate_void_fraction(self, i, j):
        """在交错网格点上插值计算空泡分数"""
        if i == 0:
            return self.scalar_grid['void_fraction'][0,j]
        elif i == self.n_axial:
            return self.scalar_grid['void_fraction'][-1,j]
        else:
            return 0.5 * (self.scalar_grid['void_fraction'][i-1,j] + 
                         self.scalar_grid['void_fraction'][i,j])
                         
    def _calculate_momentum_flux(self, k, i, j):
        """计算动量通量（使用迎风格式）"""
        if k == 'l':
            u = self.axial_grid['vel_liquid']
            alpha = 1 - self.scalar_grid['void_fraction']
            rho = self.subchannel.rho_liquid
        else:
            u = self.axial_grid['vel_vapor']
            alpha = self.scalar_grid['void_fraction']
            rho = self.subchannel.rho_vapor
            
        # 使用迎风格式计算通量
        if u[i,j] > 0:
            if i > 0:
                flux = u[i,j] * (alpha[i-1,j] * rho[i-1,j] * u[i,j])
            else:
                flux = u[i,j] * self._get_inlet_momentum_flux(k, j)
        else:
            if i < self.n_axial:
                flux = u[i,j] * (alpha[i,j] * rho[i,j] * u[i,j])
            else:
                flux = u[i,j] * self._get_outlet_momentum_flux(k, j)
                
        return flux